# -*- mode: python -*-
###
#
# (c) 2000-2007 Wijnand Modderman <wijnand at tehmaze dot com>
# All rights reserved. See docs/LICENSE
#
###

__revision__ = '$Hg: ssh.py,v 32e65e8b181f 2007/01/14 02:18:23 +0100 $'

# system imports
import os
import socket

# DDB imports
from ddb.lang import _
from ddb.log import log
from ddb.lang import lang
from ddb.transport import ITransport, ITransportProtocol
from ddb.user import user, userRegister
from ddb.event import event
from ddb.settings import *

# Twisted imports
from twisted.cred import portal, checkers
from twisted.conch import error, avatar
from twisted.conch.checkers import SSHPublicKeyDatabase
from twisted.conch.ssh import factory, userauth, connection, keys, session
from twisted.internet import reactor, protocol, defer
from zope.interface import implements
from twisted.python import components
from twisted.protocols import basic

class sshProtocol(protocol.Protocol, ITransportProtocol):
    type      = 'ssh'
    usePrefix = False
    multiLine = True

    delimiter = '\n'
    prompt    = '$u@$h$w> '
    host      = socket.getfqdn()
    cwd       = []
    BACKSPACE = ''.join(map(chr, [8, 32, 8]))

    def __init__(self):
        ITransportProtocol.__init__(self)
        self._esc = False

    #
    # Raw stream i/o
    #

    def dataReceived(self, data):
        """Called if we receive a chunk of data."""
        self.dataReceiveFilter(data)
        bufferLength = len(self.buffer)
        if bufferLength:
            if bufferLength > 512:
                # Flood! Goodbye!
                self.loseConnection()
            else:
                self.dataParse()

    def dataReceiveFilter(self, data):
        """Filter all control characters."""
        for char in data:
            charCode = ord(char)
            #print "<<< ssh in", charCode, CHARMAP.has_key(charCode) and CHARMAP[charCode] or '?'
            if CHARMAP.has_key(charCode):
                hook = '_handleChar_%s' % CHARMAP[charCode]
                if hasattr(self, hook):
                    getattr(self, hook)()
                    return
                else:
                    log.msg(_('Can not handle char %02x (%s)') % (charCode, CHARMAP[charCode]), 
                        debug=True)
            self._handleChar(char)

    def dataParse(self):
        """The buffer has data, use it."""
        lines = self.buffer.split('\r')
        if len(lines) > 1:
            # There is a \r in the data
            self.buffer = ''
            for line in lines:
                if lines[-1] == '':
                    lines.pop()
                self.dataGotLine(line)

    def dataGotLine(self, line):
        """Here we should process the buffer..."""
        self.sendLine('')
        print '<<<', line
        self.handle(line)
    
    def send(self, data):
        """Send data to the console."""
        reactor.callLater(0, self.transport.write, data) # make it buffered!
        #print ">>> ssh out: ", ', '.join(map(str, map(ord, [x for x in data])))

    def sendLine(self, data):
        """Send a line to the console."""
        self.send("%s\r\n" % data)
        #self.windowMain.pushLine(data)

    #
    # Character input handling
    #

    def _handleChar(self, char):
        if self._esc: # was an escaped char
            charCode = ord(char)
            #print "<<< ssh in", charCode, CHARMAP.has_key(charCode) and CHARMAP[charCode] or '?'
            if ESCAPES.has_key(charCode):
                hook = '_handleChar_ESC_%s' % ESCAPES[charCode]
                if hasattr(self, hook):
                    getattr(self, hook)()
                else:
                    log.msg(_('Can not handle escape sequence %02x (%s)') % (charCode, ESCAPES[charCode]), 
                        debug=True)
        self.buffer = self.buffer + char
        if self.echo:
            self.send(char)

    def _handleChar_ETX(self):
        """^c received."""
        self.sendLine('')
        self.buffer = ''
        self.sendPrompt()

    def _handleChar_EOT(self):
        """^d received."""
        if not len(self.buffer):
            self.sendLine('')
            self.sendLine(lang('bye'))
            reactor.callLater(0, self.loseConnection)
        else:
            self._handle_ETX()

    def _handleChar_CR(self):
        """Carrage return received (^m)."""
        if len(self.buffer):
            self.buffer = '%s\r' % (self.buffer)

    def _handleChar_DEL(self):
        """Backspace received."""
        if len(self.buffer) > 0:
            self.send(self.BACKSPACE)
            self.buffer = self.buffer[:-1]

    def _handleChar_HT(self):
        """Tab received (^i)."""
        pass # ignore for now

    def _handleChar_FF(self):
        """^l received."""
        self.sendLine('')
        self.sendPrompt()

    def _handleChar_ETB(self):
        """^w received."""
        if self.buffer:
            pos = 0
            for i in range(len(self.buffer), 0, -1):
                j = i-1
                if self.buffer[j] == ' ':
                    pos = j
                    break
            #self.buffer = self.buffer[0:pos]
            #self.sendPrompt()
            if pos < len(self.buffer):
                for i in range(pos, len(self.buffer)):
                    self.send(self.BACKSPACE)
                self.buffer = self.buffer[0:pos]

    def _handleChar_ESC(self):
        """Escape received (^[)."""
        self._esc = True

    def _handleChar_ESC_UP(self):
        """Arrow up."""

    #
    # Events on the stream
    #

    def connectionMade(self):
        self.buffer   = ''
        self.echo     = True
        self.identity = self.session.avatar.username
        self.userJoined(self.identity, 'ssh')
        # XXX: hack to set the user username
        userRecord = self.users.users[self.identity]
        userRecord._disableCache()
        userRecord.username = self.identity
        self.sendMotd()
        self.sendPrompt()

    def loseConnection(self):
        self.transport.loseConnection()
        #self.userLeft()

    #
    # Status updates about others
    #

    def userJoined(self, senderMask, target):
        ITransportProtocol.joinedConference(self, self.identity, senderMask, target)

    def userLeft(self, senderMask, target):
        senderIdentity = self.session.avatar.username
        ITransportProtocol.partedConference(self, self.identity, senderMask, target)

    #
    # Misc.
    #

    def sendPrompt(self):
        user = self.session.avatar.username
        if self.cwd:
            path = '(%s)' % '-'.join(cwd)
        else:
            path = ''
        self.send(self.prompt.replace('$u', user).replace('$h', self.host).replace('$w', path))
        self.send(self.buffer)

    def sendMotd(self):
        """Send the Message Of The Day to the console."""
        if os.path.isfile('config/motd'):
            for line in open('config/motd', 'r').readlines():
                self.sendLine(line.rstrip())
        self.sendLine('')
        self.sendLine(_('Welcome to DDB %s') % DDB_VERSION)
        self.sendLine('')
        self.sendLine(_('Limited terminal control supported'))
        self.sendLine('')
        self.sendLine(_('Type "help" for help'))
        self.sendLine(_('Type "apro" to search in the available commandos'))
        self.sendLine('')

    def handle(self, line):
        user = self.users.update(self.session.avatar.username, 'ssh')
        event(self, 'message', user=user, target='ssh', message=line)
        self.sendPrompt()

    def reply(self, event, msg):
        msg = self._fixFormatting(msg)
        for line in msg.split(self.delimiter):
            self.sendLine(line)

    def _fixFormatting(self, line):
        """Converts IRC formatting to ANSI."""
        chars  = {'\x02': 0}
        output = ''
        for char in line:
            if char in chars.keys():
                orgchar = char
                if char == '\x02':  # bold
                    if chars[char]%2:
                        char = '\x1b[0m'
                    else:
                        char = '\x1b[1m'
                chars[orgchar] += 1
            output = '%s%s' % (output, char)
        return output

class sshSession:
    def __init__(self, avatar):
        """
        We don't use it, but the adapter is passed the avatar as its first
        argument.
        """
        self.avatar = avatar

    def getPty(self, term, windowSize, attrs):
        log.msg('pty request for terminal type "%s", window size: %s' % (term, str(windowSize)), debug=True)

    def execCommand(self, proto, cmd):
        raise Exception(_('Executing commands is not allowed'))

    def openShell(self, trans):
        ep = sshProtocol()
        # some callback references
        ep.transport = trans
        ep.session   = self
        ep.makeConnection(trans)
        trans.makeConnection(session.wrapProtocol(ep))

    def eofReceived(self):
        pass

    def closed(self):
        pass

class sshAvatar(avatar.ConchUser):
    def __init__(self, username, factory):
        avatar.ConchUser.__init__(self)
        self.username = username
        self.factory  = factory
        self.channelLookup.update({'session':session.SSHSession})

class sshRealm:
    implements(portal.IRealm)

    def __init__(self, factory):
        self.factory = factory

    def requestAvatar(self, avatarId, mind, *interfaces):
        return interfaces[0], sshAvatar(avatarId, self.factory), lambda: None

class sshChecker(checkers.InMemoryUsernamePasswordDatabaseDontUse):
    def __init__(self):
        checkers.InMemoryUsernamePasswordDatabaseDontUse.__init__(self)
        self.addUser('test', 'test')

class sshKeyChecker(SSHPublicKeyDatabase):
    def checkKey(self, credentials):
        userKeys = userRegister.keyGet(credentials.username)
        # we need to convert the ascii keys here
        if userKeys:
            userKeys = [keys.getPublicKeyString(data=key) for key in userKeys]
        return [credentials.blob in userKeys]

class sshFactory(factory.SSHFactory):
    services = {
        'ssh-userauth':   userauth.SSHUserAuthServer,
        'ssh-connection': connection.SSHConnection
        }

    def __init__(self, manager, name, config):
        self.manager     = manager
        self.name        = name
        self.config      = config
        self.portal      = portal.Portal(sshRealm(self))
        # hmm do we really want password logins? need a way to let the user
        # set passwords first then
        #self.portal.registerChecker(sshChecker())
        self.portal.registerChecker(sshKeyChecker())

# register zope adapter
components.registerAdapter(sshSession, sshAvatar, session.ISession)

class transport(ITransport):
    """SSH Transport base class."""
    factory  = sshFactory

    def __init__(self, manager, name, config):
        ITransport.__init__(self, manager, name, config)
        log.msg('starting %s' % str(self.__class__))
        try:
            serverBind = config['bind'][0]
        except KeyError:
            serverBind = '0.0.0.0'
        try:
            serverKey = config['key'][0]
        except KeyError,e:
            log.msg(_('Could not start transport, missing private RSA "key" file option.'))
            return
        if not os.path.isfile(serverKey):
            log.msg(_('Could not start transport, missing private RSA key file "%s".') % serverKey)
            return
        if not os.path.isfile(serverKey + '.pub'):
            log.msg(_('Could not start transport, missing public RSA key file "%s.pub".') % serverKey)
            return
        try:
            serverPort = int(config['port'][0])
        except KeyError, e:
            log.msg(_('Could not start transport, missing configuration option "%s"') % str(e))
            return
        self.factory.publicKeys  = {'ssh-rsa': keys.getPublicKeyString(serverKey + '.pub')}
        self.factory.privateKeys = {'ssh-rsa': keys.getPrivateKeyObject(serverKey)}
        self.server = self.factory(manager, name, config)
        reactor.listenTCP(serverPort, self.server)

    def start(self):
        self.run = True

    def stop(self):
        self.run = False

#
# We are going to define some hashes here to aid several lookups
#

# For control characters, http://en.wikipedia.org/wiki/ASCII
CHARMAP = {}
CHARMAP[0x00] = 'NUL'   # Null
CHARMAP[0x01] = 'SOH'   # Start of Header
CHARMAP[0x02] = 'STX'   # Start of Text
CHARMAP[0x03] = 'ETX'   # End of Text
CHARMAP[0x04] = 'EOT'   # End of Transmission
CHARMAP[0x09] = 'HT'    # Horizontal Tab
CHARMAP[0x0a] = 'LF'    # Line Feed
CHARMAP[0x0c] = 'FF'    # Form Feed
CHARMAP[0x0d] = 'CR'    # Carriage Return
CHARMAP[0x16] = 'SYN'   # Synchronous Idle
CHARMAP[0x17] = 'ETB'   # End of Transmission Block
CHARMAP[0x1b] = 'ESC'   # Escape
CHARMAP[0x7f] = 'DEL'   # Delete
ESCAPES = {}
ESCAPES[0x00] = 'NUL'   # None?

# vim:ts=4:sw=4:et:
